package com.example.gpu;

import jcuda.*;
import jcuda.driver.*;

import static jcuda.driver.JCudaDriver.*;

public class JCudaVectorAddExample {
    
    // CUDA核函数代码
    private static final String CUDA_CODE = 
        "extern \"C\" \n" +
        "__global__ void add(int n, float *a, float *b, float *c)\n" +
        "{\n" +
        "    int i = blockIdx.x * blockDim.x + threadIdx.x;\n" +
        "    if (i < n) {\n" +
        "        c[i] = a[i] + b[i];\n" +
        "    }\n" +
        "}\n";

    public static void main(String[] args) {
        // 初始化JCuda驱动
        cuInit(0);
        
        // 获取设备数量
        int[] deviceCount = new int[1];
        cuDeviceGetCount(deviceCount);
        
        if (deviceCount[0] == 0) {
            System.out.println("未检测到NVIDIA GPU设备");
            return;
        }
        
        // 获取第一个GPU设备
        CUdevice device = new CUdevice();
        cuDeviceGet(device, 0);
        
        // 创建CUDA上下文
        CUcontext context = new CUcontext();
        cuCtxCreate(context, 0, device);
        
        // 创建模块并加载CUDA代码
        CUmodule module = new CUmodule();
        CUfunction function = new CUfunction();
        
        // 编译并加载CUDA代码
        compileAndLoadKernel(module, function);
        
        // 向量大小
        int n = 1000000;
        
        // 为输入和输出数据分配内存
        float[] a = new float[n];
        float[] b = new float[n];
        float[] c = new float[n];
        
        // 初始化输入数据
        for (int i = 0; i < n; i++) {
            a[i] = (float)i;
            b[i] = (float)i * 2.0f;
        }
        
        // 在GPU上分配内存
        CUdeviceptr d_a = new CUdeviceptr();
        CUdeviceptr d_b = new CUdeviceptr();
        CUdeviceptr d_c = new CUdeviceptr();
        
        cuMemAlloc(d_a, n * Sizeof.FLOAT);
        cuMemAlloc(d_b, n * Sizeof.FLOAT);
        cuMemAlloc(d_c, n * Sizeof.FLOAT);
        
        // 将数据从主机内存复制到设备内存
        cuMemcpyHtoD(d_a, Pointer.to(a), n * Sizeof.FLOAT);
        cuMemcpyHtoD(d_b, Pointer.to(b), n * Sizeof.FLOAT);
        
        // 设置线程块和网格大小
        int blockSize = 256;
        int gridSize = (n + blockSize - 1) / blockSize;
        
        // 准备核函数参数
        Pointer kernelParams = Pointer.to(
            Pointer.to(new int[]{n}),
            Pointer.to(d_a),
            Pointer.to(d_b),
            Pointer.to(d_c)
        );
        
        // 计时开始
        long startTime = System.currentTimeMillis();
        
        // 调用CUDA核函数
        cuLaunchKernel(
            function,
            gridSize, 1, 1,    // 网格维度
            blockSize, 1, 1,   // 块维度
            0, null,           // 共享内存大小和流
            kernelParams, null // 内核参数
        );
        
        // 确保内核执行完成
        cuCtxSynchronize();
        
        // 计时结束
        long endTime = System.currentTimeMillis();
        
        // 将结果从设备内存复制到主机内存
        cuMemcpyDtoH(Pointer.to(c), d_c, n * Sizeof.FLOAT);
        
        // 验证结果
        boolean passed = true;
        for (int i = 0; i < Math.min(10, n); i++) {
            float expected = a[i] + b[i];
            if (Math.abs(c[i] - expected) > 1e-5) {
                passed = false;
                break;
            }
        }
        
        // 打印结果
        System.out.println("GPU计算结果验证: " + (passed ? "通过" : "失败"));
        System.out.println("GPU计算耗时: " + (endTime - startTime) + " ms");
        
        // 在CPU上执行相同计算并比较性能
        startTime = System.currentTimeMillis();
        float[] cCpu = new float[n];
        for (int i = 0; i < n; i++) {
            cCpu[i] = a[i] + b[i];
        }
        endTime = System.currentTimeMillis();
        System.out.println("CPU计算耗时: " + (endTime - startTime) + " ms");
        
        // 释放GPU内存
        cuMemFree(d_a);
        cuMemFree(d_b);
        cuMemFree(d_c);
        
        // 清理资源
        cuModuleUnload(module);
        cuCtxDestroy(context);
    }
    
    /**
     * 编译并加载CUDA核函数
     */
    private static void compileAndLoadKernel(CUmodule module, CUfunction function) {
        try {
            // 对于演示目的，这里简化处理
            // 实际应用中应该使用cuModuleLoadDataEx或cuModuleLoad加载编译好的PTX或cubin文件
            // 为了避免复杂的编译过程，这里使用一个简单的占位符方法
            throw new UnsupportedOperationException("在实际运行前，需要将CUDA代码编译为PTX或cubin文件");
        } catch (Exception e) {
            System.out.println("加载CUDA核函数失败: " + e.getMessage());
            System.out.println("注意: 要运行此示例，需要:");
            System.out.println("1. 安装NVIDIA CUDA Toolkit");
            System.out.println("2. 将CUDA代码编译为PTX文件");
            System.out.println("3. 修改代码使用cuModuleLoad加载编译好的PTX文件");
        }
    }
}